/*
 * Copyright 2013 Eduard Sedov.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.es.util;

import java.util.AbstractMap;
import java.util.Arrays;
import java.util.Set;

/**
 *
 * @author Eduard Sedov
 */
public class LinearMap<K, V> extends AbstractMap<K, V> {

    /**
     * The maximum count of buckets, used if a higher value is implicitly
     * specified by either of the constructors with arguments. MUST be a power
     * of two <= 1<<30.
     */
    protected static final int MAXIMUM_BUCKETS_COUNT = 1 << 30;
    /**
     * The default count of buckets - MUST be a power of two.
     */
    protected static final int DEFAULT_INITIAL_BUCKETS_COUNT = 2;
    /**
     * The load factor used when none specified in constructor.
     */
    protected static final float DEFAULT_LOAD_FACTOR = 0.75f;
    protected static final Object NULL_KEY = new Object();
    protected int bucketSize;
    protected int initialBucketsCount;
    protected int bucketsCount;
    protected int level;
    protected int thisLevelBucketsCount;
    protected int nextBucketSplitIndex;
    protected Bucket[] buckets;
    protected int size;
    protected int capacity;
    protected float loadFactor;
    protected int threshold;
    protected int thisLevelIndexOperand;
    protected int nextLevelIndexOperand;

    public LinearMap() {
        this(DEFAULT_INITIAL_BUCKETS_COUNT, DEFAULT_LOAD_FACTOR);
    }

    public LinearMap(int bucketSize, int initialBucketsCount) {
        this(bucketSize, initialBucketsCount, DEFAULT_LOAD_FACTOR);
    }

    public LinearMap(int initialBucketsCount, float loadFactor) {
        this(4, initialBucketsCount, loadFactor);
    }

    public LinearMap(int bucketSize, int initialBucketsCount, float loadFactor) {
        bucketSize = Math.min(Math.max(1, bucketSize), MAXIMUM_BUCKETS_COUNT);
        int tmpBucketSize = 1;
        while (tmpBucketSize < bucketSize) {
            tmpBucketSize <<= 1;
        }

        initialBucketsCount = Math.min(Math.max(2, initialBucketsCount), MAXIMUM_BUCKETS_COUNT);
        int tmpBucketsCount = 1;
        while (tmpBucketsCount < initialBucketsCount) {
            tmpBucketsCount <<= 1;
        }
        this.loadFactor = loadFactor;
        init(tmpBucketSize, tmpBucketsCount);
    }

    private void init(int bucketSize, int initialBucketsCount) {
        this.bucketSize = bucketSize;
        this.initialBucketsCount = initialBucketsCount;
        this.bucketsCount = initialBucketsCount;
        this.thisLevelBucketsCount = bucketsCount;
        this.buckets = new Bucket[this.thisLevelBucketsCount];
        for (int i = 0; i < buckets.length; i++) {
            buckets[i] = new Bucket(bucketSize);
        }
        this.threshold = recalculateThreshold();
        this.nextBucketSplitIndex = 0;
        this.size = 0;
        this.level = 0;
        this.thisLevelIndexOperand = indexOperandForLevel(level);
        this.nextLevelIndexOperand = indexOperandForLevel(level + 1);
    }

    private int recalculateThreshold() {
        return (int) (this.bucketsCount * this.bucketSize * this.loadFactor);
    }

    @Override
    public V put(K key, V value) {
        return put0(key, value, false);
    }

    protected V put0(Object key, Object value, boolean rehash) {
        if (key == null) {
            key = NULL_KEY;
        }
        int hash = hash(key);
        int bucketIndex = bucketIndexFor(hash);
        Bucket<K, V> bucket = buckets[bucketIndex];

        int bucketLength = bucketSize;

        do {
            int[] hashes = bucket.hashes;
            Object[] keys = bucket.keys;
            Object[] values = bucket.values;
            int index = Arrays.binarySearch(hashes, hash);
            if (index < 0) {
                index = -index - 1;
                if (index < bucketLength && keys[bucketLength - 1] == null) {
                    System.arraycopy(hashes, index, hashes, index + 1, bucketLength - index - 1);
                    System.arraycopy(keys, index, keys, index + 1, bucketLength - index - 1);
                    System.arraycopy(values, index, values, index + 1, bucketLength - index - 1);
                    hashes[index] = hash;
                    keys[index] = key;
                    values[index] = value;
                    if (!rehash) {
                        size++;
                    }
                    return null;
                }
            } else if (key.equals(keys[index])) {
                V oldValue = (V) values[index];
                values[index] = value;
                return oldValue;
            }


//            Object[] keys = bucket.keys;
//
//            for (int i = 0; i < bucketLength; i++) {
//                Object tmpKey = keys[i];
//                if (tmpKey == null) {
//                    keys[i] = key;
//                    bucket.values[i] = value;
//                    if (!rehash) {
//                        size++;
//                    }
//                    return null;
//                } else if (tmpKey.equals(key)) {
//                    V oldValue = (V) bucket.values[i];
//                    bucket.values[i] = value;
//                    return oldValue;
//                }
//            }

            Bucket overflowBucket = bucket.overflowBucket;
            if (overflowBucket == null) {
                overflowBucket = new Bucket(bucketSize);
                overflowBucket.hashes[0] = hash;
                overflowBucket.keys[0] = key;
                overflowBucket.values[0] = value;
                bucket.overflowBucket = overflowBucket;
                if (!rehash) {
                    size++;
                    overflowOccurred();
                }
                return null;
            } else {
                bucket = bucket.overflowBucket;
            }


        } while (true);
    }

    @Override
    public V get(Object key) {
        if (key == null) {
            key = NULL_KEY;
        }
        int bucketIndex = bucketIndexFor(hash(key));

        Bucket<K, V> bucket = buckets[bucketIndex];

        int bucketLength = bucketSize;

        while (bucket != null) {
            Object[] keys = bucket.keys;
            for (int i = 0; i < bucketLength; i++) {
                if (key.equals(keys[i])) {
                    return (V) bucket.values[i];
                }
            }

            bucket = bucket.overflowBucket;
        }
        return null;
    }

    @Override
    public V remove(Object key) {
        if (key == null) {
            key = NULL_KEY;
        }
        int bucketIndex = bucketIndexFor(hash(key));

        Bucket<K, V> bucket = buckets[bucketIndex];

        int bucketLength = bucketSize;

        while (bucket != null) {
            Object[] keys = bucket.keys;
            for (int i = 0; i < bucketLength; i++) {
                if (key.equals(keys[i])) {
                    keys[i] = null;
                    size--;
                    return (V) bucket.values[i];
                }
            }

            bucket = bucket.overflowBucket;
        }
        return null;
    }

    @Override
    public void clear() {
        init(bucketSize, initialBucketsCount);
    }

    protected int hash(Object key) {
        if (key == NULL_KEY) {
            return 0;
        }
        int h = 0 ^ key.hashCode();
        h ^= (h >>> 20) ^ (h >>> 12);
        return h ^ (h >>> 7) ^ (h >>> 4);
    }

    protected int bucketIndexFor(int hash) {
        // modulo devision can be replaced by '&' if second operand is a power of two
        // x % 2 pow n == x & (2 pow n - 1)
        int index = hash & (thisLevelIndexOperand - 1);
        if (index < nextBucketSplitIndex) {
            return hash & (nextLevelIndexOperand - 1);
        }
        return index;
    }

    protected void overflowOccurred() {
        if (size > threshold) {
            Bucket[] tmp = buckets;
            buckets = new Bucket[++bucketsCount];
            System.arraycopy(tmp, 0, buckets, 0, tmp.length);
            buckets[buckets.length - 1] = new Bucket(bucketSize);
            threshold = recalculateThreshold();

            Bucket bucketToSplit = buckets[nextBucketSplitIndex];
            buckets[nextBucketSplitIndex] = new Bucket(bucketSize);
            nextBucketSplitIndex++;

            int bucketLength = bucketSize;
            while (bucketToSplit != null) {
                Object[] keys = bucketToSplit.keys;
                for (int i = 0; i < bucketLength; i++) {
                    Object key = keys[i];
//                    int bucketIndex = bucketIndexFor(hash(key));
//                    Bucket targetBucket = buckets[bucketIndex];
                    if (key != null) {
                        put0(key == NULL_KEY ? null : key, bucketToSplit.values[i], true);
                    }
                }

                bucketToSplit = bucketToSplit.overflowBucket;
            }


            if (nextBucketSplitIndex == thisLevelBucketsCount) {
                level++;
                nextBucketSplitIndex = 0;
                thisLevelBucketsCount = bucketsCount;
                thisLevelIndexOperand = nextLevelIndexOperand;
                nextLevelIndexOperand = indexOperandForLevel(level);
            }
        }
    }

    protected int indexOperandForLevel(int level) {
        return level == 0 ? initialBucketsCount : (2 << level - 1) * initialBucketsCount;
    }

    @Override
    public Set<Entry<K, V>> entrySet() {
        throw new UnsupportedOperationException("Not supported yet.");
    }

    protected class Bucket<K, V> {

        protected int[] hashes;
        protected Object[] keys;
        protected Object[] values;
        protected Bucket<K, V> overflowBucket;

        public Bucket(int bucketSize) {
            hashes = new int[bucketSize];
            Arrays.fill(hashes, Integer.MAX_VALUE);
            keys = new Object[bucketSize];
            values = new Object[bucketSize];
        }
    }

    public static void main(String[] args) throws Exception {
        LinearMap m = new LinearMap(128, 2, 0.75f);


        int count = 10000;

        for (int i = 0; i < count; i++) {
            Integer v = new Integer(i);
            try {
                m.put(v, v);
            } catch (Exception e) {
                System.err.println(i);
                throw e;
            }
        }


//        for (int i = 0; i < count; i++) {
//            Integer ii = new Integer(i);
//            if (!ii.equals(m.get(ii))) {
//                System.out.println(ii + " " + m.get(ii));
//            }
//        }

    }
}
